import sys
import pandas as pd
from lib.conn import Conn
from lib.func import *
from functools import reduce

DIR = os.path.dirname(os.path.realpath(__file__))

pd.set_option('display.max_rows', None)
pd.set_option('display.width', None)
pd.set_option('display.max_columns', None)



MIN_PERCENTILE = 70
POS_COUNT_TARGET = ''
REFERENCE_TICKER = 'SPY'
ALL_STOCKS = True
TICKER_INFO_FILE = ''
TICKER_INFO_DICT = ''

TITLE_RANK = "Rank"
TITLE_TICKER = "Ticker"
TITLE_TICKERS = "Tickers"
TITLE_INDUSTRY = "Industry"
TITLE_PERCENTILE = "Percentile"
TITLE_1M = "1 Month Ago"
TITLE_3M = "3 Months Ago"
TITLE_6M = "6 Months Ago"
TITLE_RS = "Relative Strength"

if not os.path.exists('output'):
    os.makedirs('output')


def relative_strength(closes: pd.Series, closes_ref: pd.Series):
    rs_stock = strength(closes)
    rs_ref = strength(closes_ref)
    rs = (1 + rs_stock) / (1 + rs_ref) * 100
    rs = int(rs*100) / 100
    return rs


def strength(closes: pd.Series):
    """Calculates the performance of the last year (most recent quarter is weighted double)"""
    try:
        quarters1 = quarters_perf(closes, 1)
        quarters2 = quarters_perf(closes, 2)
        quarters3 = quarters_perf(closes, 3)
        quarters4 = quarters_perf(closes, 4)
        return 0.4*quarters1 + 0.2*quarters2 + 0.2*quarters3 + 0.2*quarters4
    except:
        return 0


def quarters_perf(closes: pd.Series, n):
    length = min(len(closes), n*int(252/4))
    prices = closes.tail(length)
    pct_chg = prices.pct_change().dropna()
    perf_cum = (pct_chg + 1).cumprod() - 1
    return perf_cum.tail(1).item()


def rankings():
    """Returns a dataframe with percentile rankings for relative strength"""
    relative_strengths = []
    ranks = []
    industries = {}
    ind_ranks = []
    stock_rs = {}
    conn = Conn()
    stocks = conn.fetch_all('select ts_code,industry from td_stock_ak order by ts_code')
    for stock in stocks:
        try:
            ticker = stock['ts_code']
            closes = get_stock_close(ticker)
            closes_ref = get_index_close('000300.SH')
            industry = stock['industry']
            if len(closes) >= 6*20 and industry != "n/a" and len(industry.strip()) > 0:
                closes_series = pd.Series(closes)
                closes_ref_series = pd.Series(closes_ref)
                rs = relative_strength(closes_series, closes_ref_series)
                month = 20
                tmp_percentile = 100
                rs1m = relative_strength(closes_series.head(-1*month), closes_ref_series.head(-1*month))
                rs3m = relative_strength(closes_series.head(-3*month), closes_ref_series.head(-3*month))
                rs6m = relative_strength(closes_series.head(-6*month), closes_ref_series.head(-6*month))

                # if rs is too big assume there is faulty price data
                if rs < 600:
                    # stocks output
                    ranks.append(len(ranks)+1)
                    relative_strengths.append((0, stock['ts_code'], industry, rs, tmp_percentile, rs1m, rs3m, rs6m))
                    stock_rs[ticker] = rs
                    # industries output
                    if industry not in industries:
                        industries[industry] = {
                            "info": (0, industry, 0, 99, 1, 3, 6),
                            TITLE_RS: [],
                            TITLE_1M: [],
                            TITLE_3M: [],
                            TITLE_6M: [],
                            TITLE_TICKERS: []
                        }
                        ind_ranks.append(len(ind_ranks)+1)
                    industries[industry][TITLE_RS].append(rs)
                    industries[industry][TITLE_1M].append(rs1m)
                    industries[industry][TITLE_3M].append(rs3m)
                    industries[industry][TITLE_6M].append(rs6m)
                    industries[industry][TITLE_TICKERS].append(ticker)
        except KeyError:
            print(f'Ticker {ticker} has corrupted data.')
    dfs = []
    suffix = ''
    def getDfView(industry_entry):
        return industry_entry["info"]

    def sum(a,b):
        return a+b

    def getRsAverage(industries, industry, column):
        rs = reduce(sum, industries[industry][column])/len(industries[industry][column])
        rs = int(rs*100) / 100 # round to 2 decimals
        return rs

    def rs_for_stock(ticker):
        return stock_rs[ticker]

    def getTickers(industries, industry):
        return ",".join(sorted(industries[industry][TITLE_TICKERS], key=rs_for_stock, reverse=True))

    df = pd.DataFrame(relative_strengths, columns=[TITLE_RANK, TITLE_TICKER, TITLE_INDUSTRY, TITLE_RS, TITLE_PERCENTILE, TITLE_1M, TITLE_3M, TITLE_6M])
    df[TITLE_PERCENTILE] = pd.qcut(df[TITLE_RS], 100, labels=False, duplicates="drop")
    df[TITLE_1M] = pd.qcut(df[TITLE_1M], 100, labels=False, duplicates="drop")
    df[TITLE_3M] = pd.qcut(df[TITLE_3M], 100, labels=False, duplicates="drop")
    df[TITLE_6M] = pd.qcut(df[TITLE_6M], 100, labels=False, duplicates="drop")
    df = df.sort_values(([TITLE_RS]), ascending=False)
    df[TITLE_RANK] = ranks
    out_tickers_count = 0
    conn = Conn()
    conn.insert("update td_stock_ak set rs=0 where 1")
    for index, row in df.iterrows():
        print(row)
        conn.update_dict("td_stock_ak", {"rs": row[TITLE_PERCENTILE]}, {"ts_code": row[TITLE_TICKER]})
        if row[TITLE_PERCENTILE] >= MIN_PERCENTILE:
            out_tickers_count = out_tickers_count + 1
    df = df.head(out_tickers_count)




    # dfs.append(df)
    # # remove industries with only one stock
    # filtered_industries = filter(lambda i: len(i[TITLE_TICKERS]) > 1, list(industries.values()))
    # df_industries = pd.DataFrame(map(getDfView, filtered_industries), columns=[TITLE_RANK, TITLE_INDUSTRY, TITLE_RS, TITLE_PERCENTILE, TITLE_1M, TITLE_3M, TITLE_6M])
    # df_industries[TITLE_RS] = df_industries.apply(lambda row: getRsAverage(industries, row[TITLE_INDUSTRY], TITLE_RS), axis=1)
    # df_industries[TITLE_1M] = df_industries.apply(lambda row: getRsAverage(industries, row[TITLE_INDUSTRY], TITLE_1M), axis=1)
    # df_industries[TITLE_3M] = df_industries.apply(lambda row: getRsAverage(industries, row[TITLE_INDUSTRY], TITLE_3M), axis=1)
    # df_industries[TITLE_6M] = df_industries.apply(lambda row: getRsAverage(industries, row[TITLE_INDUSTRY], TITLE_6M), axis=1)
    # df_industries[TITLE_PERCENTILE] = pd.qcut(df_industries[TITLE_RS], 100, labels=False, duplicates="drop")
    # df_industries[TITLE_1M] = pd.qcut(df_industries[TITLE_1M], 100, labels=False, duplicates="drop")
    # df_industries[TITLE_3M] = pd.qcut(df_industries[TITLE_3M], 100, labels=False, duplicates="drop")
    # df_industries[TITLE_6M] = pd.qcut(df_industries[TITLE_6M], 100, labels=False, duplicates="drop")
    # df_industries[TITLE_TICKERS] = df_industries.apply(lambda row: getTickers(industries, row[TITLE_INDUSTRY]), axis=1)
    # df_industries = df_industries.sort_values(([TITLE_RS]), ascending=False)
    # ind_ranks = ind_ranks[:len(df_industries)]
    # df_industries[TITLE_RANK] = ind_ranks
    # df_industries.to_csv(os.path.join(DIR, "output", f'rs_industries{suffix}.csv'), index = False)
    # dfs.append(df_industries)
    return df


def get_stock_close(ts_code):
    conn = Conn()
    year = today()[:4]
    sql = '(select trade_date,close from td_market_daily_{} where ts_code=%s and trade_date>%s) ' \
          'union all ' \
          '(select trade_date,close from td_market_daily_{} where ts_code=%s) '. format(int(year) - 1, year)
    start = diff_day(-365)

    arr = conn.fetch_all(sql, [ts_code, start, ts_code])
    ret = []
    for item in arr:
        ret.append(float(item['close']))
    return ret


def get_index_close(ts_code):
    conn = Conn()
    sql = 'select trade_date,close from td_market_daily_index where ts_code=%s and trade_date>%s '
    start = diff_day(-365)
    arr = conn.fetch_all(sql, [ts_code, start])
    ret = []
    for item in arr:
        ret.append(float(item['close']))
    return ret


def main(skipEnter=False):
    ranks = rankings()
    print("***\nYour 'rs_stocks.csv' is in the output folder.\n***")
    if not skipEnter:
        input("Press Enter key to exit...")


if __name__ == "__main__":
    main()

